"""
Pinecone 入门示例：
- 使用 sklearn 的 digits(8x8) 数据集（1797 条，特征维度 64）。
- 如无索引则创建（维度=64，euclidean 距离），并上载全部或部分数据。
- 随机取一张图片进行查询，打印前 10 个近邻(id、距离、标签)，并输出最终预测。

环境变量：
- PINECONE_API_KEY：你的 Pinecone API Key（建议用环境变量，而不是写入代码）。
- PINECONE_INDEX（可选）：索引名，默认 "mnist-index"（示例里沿用该名称）。

首次运行若索引不存在，会自动创建并上传数据，耗时稍久，之后会很快。
"""

from __future__ import annotations

import os
import random
from typing import List, Tuple

import numpy as np
from PIL import Image
from sklearn import datasets

try:
    # 优先加载 .env（如已安装 python-dotenv）
    from dotenv import load_dotenv  # type: ignore

    load_dotenv()  # noqa: F401
except Exception:
    pass

from pinecone import Pinecone, ServerlessSpec  # type: ignore


INDEX_NAME = os.getenv("PINECONE_INDEX", "mnist-index")
API_KEY = os.getenv("PINECONE_API_KEY")

DIM = 64
METRIC = "euclidean"
REGION = os.getenv("PINECONE_REGION", "us-east-1")
CLOUD = os.getenv("PINECONE_CLOUD", "aws")
RECREATE = os.getenv("PINECONE_RECREATE", "0") in {"1", "true", "True"}
TOP_K = int(os.getenv("K", "11"))


def _get_vector_count(index) -> int:
    try:
        stats = index.describe_index_stats()
        # v5: {'namespaces': {'': {'vector_count': N}}}
        ns = stats.get("namespaces") if isinstance(stats, dict) else getattr(stats, "namespaces", {})
        if isinstance(ns, dict):
            total = 0
            for v in ns.values():
                if isinstance(v, dict) and "vector_count" in v:
                    total += int(v["vector_count"])
            return int(total)
        return 0
    except Exception:
        return 0


def ensure_index_and_upload(pc: Pinecone, name: str, X: np.ndarray, y: np.ndarray) -> None:
    """确保索引存在；若不存在则创建，并上传全部数据（必要时等待就绪）。
    结束后会打印已上传条数与当前索引向量总数，便于观察。
    """
    try:
        names = []
        try:
            # v5: list_indexes().names()
            names = list(pc.list_indexes().names())  # type: ignore[attr-defined]
        except Exception:
            # 兜底为 list[dict]
            lst = pc.list_indexes()
            # 可能是 list[{'name': '...'}] 或 {'indexes':[{'name':...}]}
            if isinstance(lst, dict) and "indexes" in lst:
                names = [ix.get("name") for ix in lst.get("indexes", [])]
            elif isinstance(lst, list):
                names = [getattr(ix, "name", getattr(ix, "id", None)) or ix.get("name") for ix in lst]  # type: ignore[union-attr]
    except Exception:
        names = []

    if RECREATE and name in names:
        print(f"索引 '{name}' 已存在，正在删除...")
        try:
            pc.delete_index(name)
            print(f"索引 '{name}' 已成功删除。")
        except Exception as e:
            print(f"删除索引失败: {e}")
        # 刷新列表
        try:
            names = list(pc.list_indexes().names())  # type: ignore[attr-defined]
        except Exception:
            lst = pc.list_indexes()
            if isinstance(lst, dict) and "indexes" in lst:
                names = [ix.get("name") for ix in lst.get("indexes", [])]
            elif isinstance(lst, list):
                names = [getattr(ix, "name", getattr(ix, "id", None)) or ix.get("name") for ix in lst]  # type: ignore[union-attr]

    if name not in names:
        pc.create_index(
            name=name,
            dimension=DIM,
            metric=METRIC,
            spec=ServerlessSpec(cloud=CLOUD, region=REGION),
        )

    index = pc.Index(name)

    # 判断是否需要上传：向量总数为 0 则上传
    need_upload = _get_vector_count(index) == 0

    if need_upload:
        print(f"Index '{name}' empty, uploading {len(X)} vectors...")
        uploaded = 0
        batch = []
        for i, (vec, label) in enumerate(zip(X, y)):
            batch.append({"id": f"train-{i}", "values": vec.tolist(), "metadata": {"label": int(label)}})
            if len(batch) >= 1000:
                index.upsert(vectors=batch)
                uploaded += len(batch)
                batch = []
        if batch:
            index.upsert(vectors=batch)
            uploaded += len(batch)
        print(f"Uploaded {uploaded} vectors to '{name}'. Waiting for index stats...")
        # 简单等待直至可见一部分向量
        import time
        for _ in range(30):  # 最多等 ~30s
            vc = _get_vector_count(index)
            if vc >= min(100, len(X)):
                break
            time.sleep(1.0)
        # 最终统计打印
        final_vc = _get_vector_count(index)
        print(f"Index '{name}' current vector_count: {final_vc}")


def build_demo_digit_three() -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """构造教程中的“数字3”8x8图像，以及其查询向量(0..16)。

    返回 (digit_img_uint8_8x8, digit_img_scaled16_8x8, query_vector_64)
    """
    import numpy as _np

    digit_3 = _np.array(
        [[0, 0, 255, 255, 255, 255, 0, 0],
         [0, 0, 0, 0, 0, 255, 0, 0],
         [0, 0, 0, 0, 0, 255, 0, 0],
         [0, 0, 0, 255, 255, 255, 0, 0],
         [0, 0, 0, 0, 0, 255, 0, 0],
         [0, 0, 0, 0, 0, 255, 0, 0],
         [0, 0, 0, 0, 0, 255, 0, 0],
         [0, 0, 255, 255, 255, 255, 0, 0]], dtype=_np.uint8
    )
    scaled16 = (digit_3.astype(_np.float32) / 255.0) * 16.0
    query_vec = scaled16.ravel().tolist()
    return digit_3, scaled16, _np.array(query_vec, dtype=_np.float32)


def to_digits_feature(img: Image.Image) -> np.ndarray:
    # 灰度并反色，然后缩放到 8x8，再归一到 0..16（与 sklearn digits 接近）
    img = img.convert("L")
    arr = 255.0 - np.array(img, dtype=np.float32)
    arr = Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8)).resize((8, 8), Image.BILINEAR)
    arr = np.array(arr, dtype=np.float32) / 255.0 * 16.0
    return arr.reshape(1, -1)


def main() -> None:
    if not API_KEY:
        raise SystemExit("请先设置环境变量 PINECONE_API_KEY 再运行示例。")

    pc = Pinecone(api_key=API_KEY)

    # 加载 digits 数据
    digits = datasets.load_digits()
    X: np.ndarray = digits.data.astype(np.float32)
    y: np.ndarray = digits.target.astype(int)

    ensure_index_and_upload(pc, INDEX_NAME, X, y)
    index = pc.Index(INDEX_NAME)

    # 教程演示版：使用“数字3”示例；若设置 USE_RANDOM=1 则改为随机样本
    use_random = os.getenv("USE_RANDOM", "0") in {"1", "true", "True"}
    if use_random:
        i = random.randrange(len(X))
        query_vec = X[i].tolist()
        true_label = int(y[i])
        demo_img = (X[i].reshape(8, 8) / 16.0 * 255.0).clip(0, 255).astype(np.uint8)
        title_hint = f"true: {true_label}"
    else:
        demo_uint8, _scaled16, qv = build_demo_digit_three()
        query_vec = qv.tolist()
        true_label = None
        demo_img = demo_uint8
        title_hint = "demo-3"

    res = index.query(vector=query_vec, top_k=TOP_K, include_metadata=True)
    matches = res.get("matches") if isinstance(res, dict) else getattr(res, "matches", [])

    # 打印前 10 近邻
    print("Top-10 neighbors:")
    labels: List[int] = []
    for m in matches:
        label = m.get("metadata", {}).get("label") if isinstance(m, dict) else getattr(m, "metadata", {}).get("label")
        score = m.get("score") if isinstance(m, dict) else getattr(m, "score", None)
        mid = m.get("id") if isinstance(m, dict) else getattr(m, "id", "?")
        labels.append(int(label))
        print(f"id: {mid}, distance: {score:.1f}, label: {label}")

    # 简单多数投票预测
    from collections import Counter

    pred = Counter(labels).most_common(1)[0][0] if labels else -1
    if true_label is not None:
        print(f"\nPredicted digit: {float(pred):.1f}  (true: {true_label})")
    else:
        print(f"\nPredicted digit: {float(pred):.1f}  ({title_hint})")

    # 保存/展示与教程类似的 8x8 放大图
    big = Image.fromarray(demo_img, mode="L").resize((256, 256), Image.NEAREST)
    out_path = "example_query.png"
    big.save(out_path)
    print(f"Saved example image -> {out_path}")


if __name__ == "__main__":
    main()
